import os
import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split
from FSdata.FSdataset import FSdata, collate_fn
import torch
import torch.utils.data as torchdata
from models.resnet import resnet50_cat
from FSdata.FSaug import RandomHflip, Resize, Compose,Normalize, Hflip
from utils.preprocessing import addFakeLabel
from utils.predicting import predict,gen_submission
from utils.metrics import cal_mAP
from utils.metrics import confusion_report
from models.drn_test import drn_d_54_cat as drn_d_54_cat_test
from utils.train_catsoftmax import train, trainlog
from FSdata.FSaug import *
from FSdata.FSdataset import FSdata, collate_fn,attr2length_map, idx2attr_map

class FSAug(object):
    def __init__(self):
        self.augment = Compose([
            Resize(size=(336,336), select=[0,1,2,5,6,7]),
            RandomUpperCrop(size=(336,336), select=[3,4]),
            RandomHflip(),
            Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ])

    def __call__(self, image,attr_idx):
        return self.augment(image,attr_idx)

class FSAugVal(object):
    def __init__(self):
        self.augment = Compose([
            Resize(size=(336, 336), select=[0,1,2,5,6,7]),
            UpperCrop(size=(336, 336), select=[3,4]),
            Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ])

    def __call__(self, image,attr_idx):
        return self.augment(image,attr_idx)



os.environ["CUDA_VISIBLE_DEVICES"] = "1"


rawdata_root = '/media/gserver/data/FashionAI'
all_pd = pd.read_csv(os.path.join(rawdata_root, 'base/Annotations/label.csv'),
                       header=None, names=['ImageName', 'AttrKey', 'AttrValues'])
all_pd['ImageName'] = all_pd['ImageName'].apply(lambda x: os.path.join('base',x))


train_pd, val_pd = train_test_split(all_pd, test_size=0.1, random_state=37,
                                    stratify=all_pd['AttrKey'])


# select part
select_AttrIdx = range(8)
select_AttrKey = [idx2attr_map[x] for x in select_AttrIdx]
train_pd = train_pd[train_pd['AttrKey'].apply(lambda x: True if x in select_AttrKey else False)]
val_pd = val_pd[val_pd['AttrKey'].apply(lambda x: True if x in select_AttrKey else False)]



data_set = {}
data_set['train'] = FSdata(root_path=rawdata_root,
                           anno_pd=train_pd,
                           transforms=FSAug(),
                           select=select_AttrIdx,
                           )
data_set['val'] = FSdata(root_path=rawdata_root,
                           anno_pd=val_pd,
                           transforms=FSAugVal(),
                         select=select_AttrIdx
                           )

data_loader = {}
data_loader['train'] = torchdata.DataLoader(data_set['train'], 8, num_workers=4,
                                            shuffle=True, pin_memory=True,collate_fn=collate_fn)
data_loader['val'] = torchdata.DataLoader(data_set['val'], batch_size=8, num_workers=4,
                                          shuffle=False, pin_memory=True,collate_fn=collate_fn)
#

# model prepare
model_name = 'resnet50_cat_bs48_smooth-[0.9674]-aug2'
resume = '/media/gserver/models/FashionAI/DRN_54d_cat_test[all]-1523063557/weights-27-2296-[0.8992]-[0.9806].pth'
model = drn_d_54_cat_test(pretrained=True, num_classes=[attr2length_map[x] for x in select_AttrIdx])

model = torch.nn.DataParallel(model)
print('resuming finetune from %s'%resume)
model.load_state_dict(torch.load(resume))
model = model.cuda()




val_preds, val_labels_str, val_attr, val_true= predict(model, data_set['val'], data_loader['val'])


val_mAP,APs, accs = cal_mAP(val_labels_str,val_preds,val_attr, data_set['val'].catidx_map)

wrong_dicts, c_mats = confusion_report(val_true,val_preds,val_attr,np.array(data_set['val'].paths), data_set['val'].catidx_map)

print wrong_dicts.keys()
print c_mats.keys()

print('=='*20)
print('val-mAP: %.4f'% (val_mAP))
for key in APs.keys():
    print('acc: %.4f, AP: %.4f %s' % (accs[key], APs[key], key))


for attr in c_mats.keys():
    print '=='*20
    print attr
    print c_mats[attr]


# import shutil
# for attr in wrong_dicts.keys():
#     for missclass in wrong_dicts[attr].keys():
#         if not os.path.exists('./missclassify/%s/%s'%(attr,missclass)):
#             os.makedirs('./missclassify/%s/%s'%(attr,missclass))
#         for img_path in wrong_dicts[attr][missclass]:
#             img_path = os.path.join(train_root,img_path)
#             assert os.path.exists(img_path)
#             img_name = img_path.split('/')[-1]
#             new_path = os.path.join('./missclassify/%s/%s'%(attr,missclass), img_name)
#             shutil.copyfile(img_path,new_path)



